#################################################################
# imports, etc.

import sys
import math
from scipy.stats import beta
from utils import *

##################################################
# load data

print('\tloading simulation data')

suff=''
actual=None
probs=None

PS_suff = ['_PS_1980_temptpu','_PS_1990_temptpu']
bounds_suff = ['_ci_lower','_ci_upper']
xi_suff = ['_lo_xi','_hi_xi','_lo_rhoxi','_hi_rhoxi']
models_suff = ['_calvo','_one_sector']

gaps = pd.read_stata(dtapath + 'inputs_calibration_gapelast.dta')
gaps=gaps.fillna(0)
actual = gaps.temp_b

probs = np.genfromtxt(inpath + 'tpuprobs_baseline.txt')
probs_perm = np.genfromtxt(inpath + 'tpuprobs_permtpu.txt')
probs_PS = [np.genfromtxt(inpath + 'tpuprobs'+s+'.txt') for s in PS_suff]
probs_bounds = [np.genfromtxt(inpath + 'tpuprobs'+s+'.txt') for s in bounds_suff]
probs_xi = [np.genfromtxt(inpath + 'tpuprobs'+s+'.txt') for s in xi_suff]
probs_models = [np.genfromtxt(inpath + 'tpuprobs'+s+'.txt') for s in models_suff]

rho_xi = probs[-1]
probs=probs[:-1]

# load data and aggregate it
df = [load_data(1),load_data(3)]
df_perm = load_data(3,'_permtpu')
df_PS = [load_data(3,'',s) for s in PS_suff]
df_bounds = [load_data(3,'',s) for s in bounds_suff]
df_xi1 = [load_data(1,'',s) for s in xi_suff]
df_xi3 = [load_data(3,'',s) for s in xi_suff]
df_models1 = [load_data(1,'',s) for s in models_suff]
df_models3 = [load_data(3,'',s) for s in models_suff]

#######################################
# annual gap coefficients

print('\tRunning regressions')

# main figures in paper
effects = [gap_regression_AKKRS(df_) for df_ in df]
effects_perm = gap_regression_AKKRS(df_perm)
effects_PS = [gap_regression_AKKRS(df_) for df_ in df_PS]
effects_bounds = [gap_regression_AKKRS(df_) for df_ in df_bounds]
effects_xi1 = [gap_regression_AKKRS(df_) for df_ in df_xi1]
effects_xi3 = [gap_regression_AKKRS(df_) for df_ in df_xi3]
effects_models1 = [gap_regression_AKKRS(df_) for df_ in df_models1]
effects_models3 = [gap_regression_AKKRS(df_) for df_ in df_models3]

# PS coefficients
effectPS_baseline = [gap_regression_PS(df_,True) for df_ in df]
effectPS_PS = [gap_regression_PS(df_,True) for df_ in df_PS]

#############################################################################
print('\tmaking plots')

# ----------------------------------
# coefficients (bench + no TPU)

fig,ax = paper_fig(True,True,True,False)

years=range(1974,2009)
ax.plot(years,actual,color='black',
        marker='o',alpha=0.7,label='Data',lw=0)

ax.plot(years,effects[1],color=colors[0],
        alpha=0.7,label='Model: with TPU',ls='-',zorder=5,lw=1.5)

ax.plot(years,effects[0],color=colors[1],
        alpha=0.7,label='Model: no TPU',linestyle='--',lw=1,zorder=6)

ax.set_yticks(range(-15,1,5))

ax.legend(frameon=True, loc='lower right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_gap_coeffs.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_gap_coeffs.csv')
plt.close('all')

# ----------------------------------
# coefficients (bench + const TPU)
fig,ax = paper_fig(True,True,True,True)

years=range(1974,2009)
ax.plot(years,actual,color='black',
        marker='o',alpha=0.7,label='Data',lw=0)

ax.plot(years,effects[1],color=colors[0],
        alpha=0.7,label='Benchmark: with TPU',ls='-',zorder=5,lw=1.5)

ax.plot(years,effects_PS[0],color=colors[2],
        alpha=0.7,label='Const. TPU from 1980',ls='-',marker='x',lw=0.75)

ax.plot(years,effects_PS[1],color=colors[3],
        alpha=0.7,label='Const. TPU from 1990',ls='-',marker='+',lw=0.75)

ax.set_yticks(range(-15,1,5))

ax.legend(frameon=True, loc='lower right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_gap_coeffs_const_tpu.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_gap_coeffs_const_tpu.csv')
plt.close('all')

# ----------------------------------
# coefficients (bench + fast adj (no TPU)
fig,ax = paper_fig(True,True,True,False)

years=range(1974,2009)
ax.plot(years,actual,color='black',
        marker='o',alpha=0.7,label='Data',lw=0)

ax.plot(years,effects[1],color=colors[0],
        alpha=0.7,label='Both models: with TPU',ls='-',zorder=5,lw=1.5)

ax.plot(years,effects[0],color=colors[1],
        alpha=0.7,label='Benchmark: no TPU',linestyle='--',lw=1,zorder=6)

ax.plot(years,effects_models1[0],color=colors[2],
        alpha=0.7,label='Fast adj: no TPU',ls='--',marker='x',lw=0.75)

ax.set_yticks(range(-15,1,5))

ax.legend(frameon=True, loc='lower right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_gap_coeffs_fast_adj.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_gap_coeffs_fast_adj.csv')
plt.close('all')

# ----------------------------------
# probabilities

#t = range(1973,2008)
t = range(1980,2008)

p1 = probs.copy()
p1[NR-1-1973:]=p1[NR-1-1973]
p2 = probs.copy()
p2[0:(NR-1973)] = probs[(NR-1973)]

p2_PS = (probs_PS[0]+0)*np.ones(len(p2))
p2_PS[0:7] = 0
p2_PS[-7:]=0

p2_PS_2 = (probs_PS[1]+0)*np.ones(len(p2))
p2_PS_2[0:17] = 0
p2_PS_2[-7:]=0

p1_calvo = probs_models[0].copy()
p1_calvo[NR-1-1973:]=p1_calvo[NR-1-1973]
p2_calvo = probs_models[0].copy()
p2_calvo[0:(NR-1973)] = probs_models[0][(NR-1973)]

fig,ax = paper_fig(False,True,True,False)
ax.plot(t,p2[7:],color=colors[0],alpha=0.7,label=r'MFN to NNTR',zorder=6,lw=1.5)
ax.bar([1975],[p1[0]],color=colors[1],width=3,label='NNTR to MFN',alpha=0.7)
ax.set_yticks(np.arange(0,1.0,0.2))
ax.set_xlim(1972,2007)
ax.set_xticks([1975, 1980, 1985, 1990, 1995, 2000, 2005])
ax.set_xticklabels(['1970s','1980','1985','1990','1995','2000','2005'],minor=False)
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')
plt.savefig(figpath + 'model_probs.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_probs.csv')
plt.close('all')

fig,ax = paper_fig(False,True,True,True)
ax.plot(t,p2[7:],color=colors[0],alpha=0.7,label=r'Benchmark',zorder=6,lw=1.5)
ax.plot(t,p2_PS[7:],color=colors[2],alpha=0.7,label="Const. TPU from 1980",linestyle='-',marker='x',lw=0.75)
ax.plot(t,p2_PS_2[7:],color=colors[3],alpha=0.7,label="Const. TPU from 1990",linestyle='-',marker='+',lw=0.75)
ax.set_yticks(np.arange(0,1.0,0.2))
ax.set_xlim(1980,2008)
ax.set_xticks([1980,1985,1990,1995,2000,2005])
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')
plt.savefig(figpath + 'model_probs_const_tpu.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_probs_const_tpu.csv')
plt.close('all')

fig,ax = paper_fig(False,True,True,False)
ax.plot(t,p2[7:],color=colors[0],alpha=0.7,label=r'Benchmark',zorder=6,lw=1.5)
ax.plot(t,p2_calvo[7:],color=colors[2],alpha=0.7,label=r'Fast adj',marker='x',lw=0.75)
ax.set_yticks(np.arange(0,1.2,0.2))
ax.set_xlim(1980,2008)
ax.set_xticks([1980,1985,1990,1995,2000,2005])
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')
plt.savefig(figpath + 'model_probs_fast_adj.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_probs_fast_adj.csv')
plt.close('all')

# ----------------------------------
# trade

def plot_agg_trade(df,ax,color,marker,ls,lab,lw,zorder,t0=0):
    df2 = df[(df.y >=1974)&(df.y<=2008)].groupby(['y'])['exports'].sum().reset_index()
    ax.plot(years[t0:],np.log(df2.exports/df2.exports[0])[t0:],color=color,linestyle=ls,marker=marker,label=lab,alpha=0.7,lw=lw,zorder=zorder)


fig,ax = paper_fig(False,True,True,False)
plot_agg_trade(df[1],ax,colors[0],None,'-','Model',lw=1.5,zorder=5)
plot_agg_trade(df[0],ax,colors[1],None,'--','No TPU',lw=1,zorder=6)
ax.legend(frameon=True, loc='lower right', fontsize=8, framealpha=1, edgecolor='white')
plt.savefig(figpath + 'model_agg_trade.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_agg_trade.csv')
plt.close('all')
    
# ----------------------------------
# PV tariff

def plot_pv_tau(df,ax,color,marker,ls,lab,flag,lw,zorder,t0=0):
    df2 = df[(df.y >=1974)&(df.y<=2008)].groupby(['y'])[['pv_tau','tau_applied']].mean().reset_index()

    if(flag):
        ax.plot(years[t0:],df2.tau_applied[t0:],color=color,linestyle=ls,marker=marker,label=lab,alpha=0.7,lw=lw)
    else:
        ax.plot(years[t0:],df2.pv_tau[t0:]-1,color=color,linestyle=ls,marker=marker,label=lab,alpha=0.7,lw=lw)

fig,ax = paper_fig(False,True,True,False)
plot_pv_tau(df[1],ax,'black','o','-','Applied tariff',1,0,zorder=1)
plot_pv_tau(df[1],ax,colors[0],None,'-','Discounted: with TPU',0,1.5,zorder=5)
plot_pv_tau(df[0],ax,colors[1],None,'--','Discounted: no TPU',0,1,zorder=6)
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')
plt.savefig(figpath + 'model_pv_tau.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_pv_tau.csv')
plt.close('all')

#-------------------------------------------------------
# bayesian stuff

# prior in 1970
priors0 = [(1,1),(2,2),(8,8),(8,4),(0.1,0.1)]
mu0 = [[] for p in priors0]
for i in range(len(priors0)):
    a = priors0[i][0]
    b = priors0[i][1]
    
    # 1970 prior + posteriors during 1971-1979
    for t in range(0,9):
        mu0[i].append(float(beta.stats(t+a,b,moments='m')))

mu1_model = probs[8:]

x = mu1_model[0]
# a/(a+b)=x --> a = a*x + b*x --> a(1-x) = b*x --> a = b * x/(1-x)
priors1 = [(b*x/(1-x),b) for b in np.linspace(1,5,5)]
mu1 = [[] for p in priors1]
for i in range(len(priors1)):

    a = priors1[i][0]
    b = priors1[i][1]

    # posteriors during 1981-2008
    for t in range(0,27):   
        mu1[i].append(float(beta.stats(a,t+b,moments='m')))

theta = np.linspace(0,1,24)

fig1,axes1 = paper_fig(False,True,True,False)
fig2,axes2 = paper_fig(False,True,True,False)

markers=[None,None,'x','+','1']
linestyles=['-','--','-','-','-']

lns2=[]
lns1=[]

lns2 = lns2 + axes2.plot(range(1981,2008),mu1_model,color='black',linewidth=1,marker='o',label='Model')
for i in range(len(priors1)):
    a = priors1[i][0]
    b = priors1[i][1]
    lns1 = lns1 + axes1.plot(theta,beta.pdf(theta,a,b),linestyle=linestyles[i],label='b=%0.1f'%(b),alpha=alpha,linewidth=0.75,color=colors[i],marker=markers[i])
    lns2 = lns2 + axes2.plot(range(1981,2008),mu1[i],linestyle=linestyles[i],label='b=%0.1f'%(b),alpha=alpha,linewidth=0.75,color=colors[i],marker=markers[i])

axes1.set_xlim(0,1)
axes2.set_xlim(1981,2008)
axes1.legend(frameon=True, framealpha=1, edgecolor='white',loc='best',prop={'size':8})
axes2.legend(frameon=True, framealpha=1, edgecolor='white',loc='best',prop={'size':8})
fig1.subplots_adjust(hspace=0.2,wspace=0.25)
fig2.subplots_adjust(hspace=0.2,wspace=0.25)

plt.sca(axes1)
plt.savefig(figpath + 'bayesian_probs_priors.pdf',bbox_inches='tight')
extract_fig_data(axes1,figpath + 'bayesian_probs_priors.csv')

plt.sca(axes2)
plt.savefig(figpath + 'bayesian_probs_posteriors.pdf',bbox_inches='tight')
extract_fig_data(axes2,figpath + 'bayesian_probs_posteriors.csv')
plt.close('all')

##############################################

columns=['f0','f1','xi','sig']
params = pd.read_csv(inpath + 'exporter_params_baseline.csv',sep=' ',header=None,names=columns)

columns=['sector_num','delast','partic','exit','inc_prem','cv']
sectors = pd.read_csv('../model/input/inputs_calibration_sector.csv',sep=' ',header=None,names=columns)
params['elast'] = sectors.delast

sector_names = ['Food, beverage, tobacco',
                'Textile, clothing, footwear',
                'Wood and straw products',
                'Paper, printing products',
                'Energy products, chemicals',
                'Rubber, plastic products',
                'Non-metallic mineral products',
                'Base metal manuf.',
                'Calendered metal manuf.',
                'Other machinery, equipment',
                'Computer, electronic, optical',
                'Electrical equipment manuf.',
                'Vehicle manuf.',
                'Furniture, other manuf.',
                'Non-manufacturing']

file = open(texpath + 'sector_params.tex','w')
file2 = open(path + '../30 Final output files/table2.tex','w')

file.write('\\begin{table}[tbh]\n\\scriptsize\n\\begin{center}\n\\caption{Sector-level model parameters}\n\\label{tab:cal_sectors}\n\\begin{tabular*}{4.25in}{l@{\\extracolsep{\\fill}}lccccc}\n\\toprule\n\\multicolumn{2}{c}{Sector} & $\\theta_g$ & $f_{g0}$ & $f_{g1}$ & $\\xi_{gH}$ & $\\sigma_{gz}$\\\\\n\\midrule\n')
file2.write('\\begin{table}[tbh]\n\\scriptsize\n\\begin{center}\n\\caption{Sector-level model parameters}\n\\label{tab:cal_sectors}\n\\begin{tabular*}{4.25in}{l@{\\extracolsep{\\fill}}lccccc}\n\\toprule\n\\multicolumn{2}{c}{Sector} & $\\theta_g$ & $f_{g0}$ & $f_{g1}$ & $\\xi_{gH}$ & $\\sigma_{gz}$\\\\\n\\midrule\n')

for i in range(len(params)):
    file.write('%d & %s & %0.2f & %0.2f & %0.2f & %0.2f & %0.2f\\\\\n' % (i+1,sector_names[i],params.elast[i],params.f0[i],params.f1[i],params.xi[i],params.sig[i]))
    #file2.write('%d & %s & %0.2f & %0.2f & %0.2f & %0.2f & %0.2f\\\\\n' % (i+1,sector_names[i],params.elast[i],params.f0[i],params.f1[i],params.xi[i],params.sig[i]))

file.write('\\bottomrule\n\\end{tabular*}\n\\end{center}\n\\end{table}\n')
file2.write('\\bottomrule\n\\end{tabular*}\n\\end{center}\n\\end{table}\n')

file.close()
file2.close()

file = open(texpath + 'other_model_numbers.txt','w')
file.write('rho_xi = %0.4f\n'%rho_xi)
file.write('PS regression coefficient for no-TPU model = %0.4f\n'%effectPS_baseline[0][0])
file.close()

file2 = open(path + '../30 Final output files/other_model_numbers.txt','w')
file2.write('rho_xi = %0.4f\n'%rho_xi)
file2.write('PS regression coefficient for no-TPU model = %0.4f\n'%effectPS_baseline[0][0])
file2.close()


##############################################
# sensitivity analyses for the appendix

# ----------------------------------
# coefficients w/ bounds

fig,ax = paper_fig(True,True,True,False)

years=range(1974,2009)
ax.plot(years,actual,color=colors[2],
        marker='o',markersize=4,alpha=0.8,label='Data')

ax.plot(years,effects[1],color=colors[1],
        alpha=0.8,label='Model: with TPU',ls='-')

ax.plot(years,effects_bounds[0],color=colors[1],
        alpha=0.8,ls='-',lw=0.5)

ax.plot(years,effects_bounds[1],color=colors[1],
        alpha=0.8,ls='-',lw=0.5)

ax.plot(years,effects[0],color=colors[0],
        alpha=0.8,label='Model: no TPU',linestyle='--')

ax.set_yticks(range(-15,1,5))

ax.legend(frameon=False, fontsize=8)

plt.savefig(figpath + 'model_gap_coeffs_bounds.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_gap_coeffs_bounds.csv')
plt.close('all')

# ----------------------------------
# probabilities w/ bounds


fig,ax = paper_fig(False,True,True,False)

#t = range(1973,2008
t = range(1980,2008)

p1 = probs.copy()
p1[NR-1-1973:]=p1[NR-1-1973]
p2 = probs.copy()
p2[0:(NR-1973)] = probs[(NR-1973)]

p1_lower = probs_bounds[0].copy()
p1_lower[NR-1-1973:]=p1_lower[NR-1-1973]
p2_lower = probs_bounds[0].copy()
p2_lower[0:(NR-1973)] = probs_bounds[0][(NR-1973)]

p1_upper = probs_bounds[1].copy()
p1_upper[NR-1-1973:]=p1_upper[NR-1-1973]
p2_upper = probs_bounds[1].copy()
p2_upper[0:(NR-1973)] = probs_bounds[1][(NR-1973)]

ax.plot(t,p2[7:],color=colors[0],alpha=0.7,label=r'MFN to NNTR')
ax.plot(t,p2_lower[7:],color=colors[0],alpha=0.7,linestyle='-',lw=0.5)
ax.plot(t,p2_upper[7:],color=colors[0],alpha=0.7,linestyle='-',lw=0.5)
ax.bar([1975],[p1[0]],color=colors[1],label='NNTR to MFN',alpha=0.7, width=3)
ax.errorbar(x=[1975],y=[p1[0]],yerr=[[p1_lower[0]],[p1_upper[0]]],
            color=colors[1],alpha=0.7,elinewidth=0.5,capsize=6,capthick=0.5)

ax.set_yticks(np.arange(0,1.0,0.2))
ax.set_xlim(1972,2007)
ax.set_xticks([1975, 1980, 1985, 1990, 1995, 2000, 2005])
ax.set_xticklabels(['1970s','1980','1985','1990','1995','2000','2005'],minor=False)
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_probs_bounds.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_probs_bounds.csv')
plt.close('all')

# ----------------------------------
# probabilities w permanent TPU version

fig,ax = paper_fig(False,True,True,False)

t = range(1980,2008)

p1 = probs.copy()
p1[NR-1-1973:]=p1[NR-1-1973]
p2 = probs.copy()
p2[0:(NR-1973)] = probs[(NR-1973)]

p1_perm = probs_perm.copy()
p1_perm[NR-1-1973:]=p1_perm[NR-1-1973]
p2_perm = probs_perm.copy()
p2_perm[0:(NR-1973)] = probs_perm[(NR-1973)]

p2_PS = probs_PS[0]+0

#ax.plot(t,p1,color=colors[0],alpha=0.7,label=r'NNTR to MFN',linestyle='--')
ax.bar([1974],[p1[0]],color=colors[1],label='Benchmark: NNTR to MFN, ',alpha=0.7, width=1.5)
ax.bar([1976],[p1_perm[0]],color=colors[3],label='Surprises: NNTR to MFN',alpha=0.7, width=1.5)
ax.plot(t,p2[7:],color=colors[0],alpha=0.7,label=r'Benchmark: MFN to NNTR')
ax.plot(t,p2_perm[7:],color=colors[2],alpha=0.7,label=r'Surprises: MFN to NNTR',marker='+',markersize=4)

ax.set_yticks(np.arange(0,1.0,0.2))
ax.set_xlim(1972,2007)
ax.set_xticks([1975, 1980, 1985, 1990, 1995, 2000, 2005])
ax.set_xticklabels(['1970s','1980','1985','1990','1995','2000','2005'],minor=False)
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_probs_sens_permtpu.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_probs_sens_permtpu.csv')
plt.close('all')

# ----------------------------------
# coefficients (bench + one sector
fig,ax = paper_fig(True,True,True,False)

years=range(1974,2009)
ax.plot(years,actual,color='black',
        marker='o',alpha=0.7,label='Data',lw=0)

ax.plot(years,effects[1],color=colors[0],
        alpha=0.7,label='Both models: with TPU',ls='-',zorder=5,lw=1.5)

ax.plot(years,effects[0],color=colors[1],
        alpha=0.7,label='Benchark: no TPU',linestyle='--',lw=1,zorder=6)

ax.plot(years,effects_models1[1],color=colors[2],
        alpha=0.7,label='One sector: no TPU',ls='--',marker='x',lw=0.75)

ax.set_yticks(range(-15,1,5))

ax.legend(frameon=True, loc='lower right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_gap_coeffs_sens_one_sector.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_gap_coeffs_sens_one_sector.csv')
plt.close('all')

# ----------------------------------
# probabilities w one sector version

fig,ax = paper_fig(False,True,True,False)

t = range(1980,2008)

p1 = probs.copy()
p1[NR-1-1973:]=p1[NR-1-1973]
p2 = probs.copy()
p2[0:(NR-1973)] = probs[(NR-1973)]

p1_os = probs_models[1].copy()
p1_os[NR-1-1973:]=probs_models[1][NR-1-1973]
p2_os = probs_models[1].copy()
p2_os[0:(NR-1973)] = probs_models[1][(NR-1973)]

p2_PS = probs_PS[0]+0

#ax.plot(t,p1,color=colors[0],alpha=0.7,label=r'NNTR to MFN',linestyle='--')
ax.bar([1974],[p1[0]],color=colors[1],label='Benchmark: NNTR to MFN',alpha=0.7, width=1.5)
ax.bar([1976],[p1_os[0]],color=colors[3],label='One sector: NNTR to MFN',alpha=0.7, width=1.5)
ax.plot(t,p2[7:],color=colors[0],alpha=0.7,label=r'Benchmark: MFN to NNTR')
ax.plot(t,p2_os[7:],color=colors[2],alpha=0.7,label=r'One sector: MFN to NNTR',marker='x',markersize=4)

ax.set_yticks(np.arange(0,1.0,0.2))
ax.set_xlim(1972,2007)
ax.set_xticks([1975, 1980, 1985, 1990, 1995, 2000, 2005])
ax.set_xticklabels(['1970s','1980','1985','1990','1995','2000','2005'],minor=False)
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_probs_sens_one_sector.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_probs_sens_one_sector.csv')
plt.close('all')

# ----------------------------------
# coefficients w/ bounds

fig,ax = paper_fig(True,True,True,False)

years=range(1974,2009)
ax.plot(years,actual,color=colors[2],
        marker='o',markersize=4,alpha=0.8,label='Data')

ax.plot(years,effects[1],color=colors[1],
        alpha=0.8,label='Model: with TPU',ls='-')

ax.plot(years,effects_bounds[0],color=colors[1],
        alpha=0.8,ls='-',lw=0.5)

ax.plot(years,effects_bounds[1],color=colors[1],
        alpha=0.8,ls='-',lw=0.5)

ax.plot(years,effects[0],color=colors[0],
        alpha=0.8,label='Model: no TPU',linestyle='--')

ax.set_yticks(range(-15,1,5))

ax.legend(frameon=False, fontsize=8)

plt.savefig(figpath + 'model_gap_coeffs_bounds.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_gap_coeffs_bounds.csv')
plt.close('all')

# ----------------------------------
# coefficients w lo vs hih xi
fig,ax = paper_fig(True,True,True,False)

years=range(1974,2009)
ax.plot(years,actual,color='black',
        marker='o',alpha=0.7,label='Data',lw=0)

ax.plot(years,effects[1],color=colors[0],
        alpha=0.7,label='All models: with TPU',ls='-',zorder=5,lw=1.5)

ax.plot(years,effects[0],color=colors[1],
        alpha=0.7,label='Benchark: no TPU',linestyle='--',lw=1,zorder=6)

ax.plot(years,effects_xi1[0],color=colors[2],
        alpha=0.7,label='Low $\\xi$: no TPU',ls='--',marker='x',lw=0.75)

ax.plot(years,effects_xi1[1],color=colors[3],
        alpha=0.7,label='High $\\xi$: no TPU',ls='--',marker='+',lw=0.75)

ax.set_yticks(range(-15,1,5))

ax.legend(frameon=True, loc='lower right', fontsize=6, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_gap_coeffs_sens_xi.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_gap_coeffs_sens_xi.csv')
plt.close('all')

# ----------------------------------
# probabilities w lo vs hi xi

fig,ax = paper_fig(False,True,True,False)

t = range(1980,2008)

p1 = probs.copy()
p1[NR-1-1973:]=p1[NR-1-1973]
p2 = probs.copy()
p2[0:(NR-1973)] = probs[(NR-1973)]

p1_lx = probs_xi[0].copy()
p1_lx[NR-1-1973:]=probs_xi[0][NR-1-1973]
p2_lx = probs_xi[0].copy()
p2_lx[0:(NR-1973)] = probs_xi[0][(NR-1973)]

p1_hx = probs_xi[1].copy()
p1_hx[NR-1-1973:]=probs_xi[1][NR-1-1973]
p2_hx = probs_xi[1].copy()
p2_hx[0:(NR-1973)] = probs_xi[1][(NR-1973)]

#ax.plot(t,p1,color=colors[0],alpha=0.7,label=r'NNTR to MFN',linestyle='--')
ax.bar([1974],[p1[0]],color=colors[1],label='Benchmark: NNTR to MFN',alpha=0.7, width=1)
ax.bar([1975],[p1_lx[0]],color=colors[3],label=r'Low $\xi$: NNTR to MFN',alpha=0.7, width=1)
ax.bar([1976],[p1_hx[0]],color=colors[5],label=r'High $\xi$: NNTR to MFN',alpha=0.7, width=1)

ax.plot(t,p2[7:],color=colors[0],alpha=0.7,label=r'Benchmark: MFN to NNTR')
ax.plot(t,p2_lx[7:],color=colors[2],alpha=0.7,label=r'Low $\xi$: MFN to NNTR',marker='x',markersize=4)
ax.plot(t,p2_hx[7:],color=colors[4],alpha=0.7,label=r'High $\xi$: MFN to NNTR',marker='+',markersize=4)

ax.set_yticks(np.arange(0,1.0,0.2))
ax.set_xlim(1972,2007)
ax.set_xticks([1975, 1980, 1985, 1990, 1995, 2000, 2005])
ax.set_xticklabels(['1970s','1980','1985','1990','1995','2000','2005'],minor=False)
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_probs_sens_xi.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_probs_sens_xi.csv')
plt.close('all')

# ----------------------------------
# coefficients w lo vs hi rho-xi
fig,ax = paper_fig(True,True,True,False)

years=range(1974,2009)
ax.plot(years,actual,color='black',
        marker='o',alpha=0.7,label='Data',lw=0)

ax.plot(years,effects[1],color=colors[0],
        alpha=0.7,label='All models: with TPU',ls='-',zorder=5,lw=1.5)

ax.plot(years,effects[0],color=colors[1],
        alpha=0.7,label='Benchark: no TPU',linestyle='--',lw=1,zorder=6)

ax.plot(years,effects_xi1[2],color=colors[2],
        alpha=0.7,label=r'Low $\rho_\xi$: no TPU',ls='--',marker='x',lw=0.75)

ax.plot(years,effects_xi1[3],color=colors[3],
        alpha=0.7,label=r'High $\rho_\xi$: no TPU',ls='--',marker='+',lw=0.75)

ax.set_yticks(range(-15,1,5))

ax.legend(frameon=True, loc='lower right', fontsize=6, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_gap_coeffs_sens_rho_xi.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_gap_coeffs_sens_rho_xi.csv')
plt.close('all')

# ----------------------------------
# probabilities w lo vs hi rho-xi

fig,ax = paper_fig(False,True,True,False)

t = range(1980,2008)

p1 = probs.copy()
p1[NR-1-1973:]=p1[NR-1-1973]
p2 = probs.copy()
p2[0:(NR-1973)] = probs[(NR-1973)]

p1_lrx = probs_xi[2].copy()
p1_lrx[NR-1-1973:]=probs_xi[2][NR-1-1973]
p2_lrx = probs_xi[2].copy()
p2_lrx[0:(NR-1973)] = probs_xi[2][(NR-1973)]

p1_hrx = probs_xi[3].copy()
p1_hrx[NR-1-1973:]=probs_xi[3][NR-1-1973]
p2_hrx = probs_xi[3].copy()
p2_hrx[0:(NR-1973)] = probs_xi[3][(NR-1973)]

#ax.plot(t,p1,color=colors[0],alpha=0.7,label=r'NNTR to MFN',linestyle='--')
ax.bar([1974],[p1[0]],color=colors[1],label='Benchmark: NNTR to MFN',alpha=0.7, width=1)
ax.bar([1975],[p1_lrx[0]],color=colors[3],label=r'Low $\rho_\xi$: NNTR to MFN',alpha=0.7, width=1)
ax.bar([1976],[p1_hrx[0]],color=colors[5],label=r'High $\rho_\xi$: NNTR to MFN',alpha=0.7, width=1)

ax.plot(t,p2[7:],color=colors[0],alpha=0.7,label=r'Benchmark: MFN to NNTR')
ax.plot(t,p2_lrx[7:],color=colors[2],alpha=0.7,label=r'Low $\rho_\xi$: MFN to NNTR',marker='x',markersize=4)
ax.plot(t,p2_hrx[7:],color=colors[4],alpha=0.7,label=r'High $\rho_\xi$: MFN to NNTR',marker='+',markersize=4)

ax.set_yticks(np.arange(0,1.0,0.2))
ax.set_xlim(1972,2007)
ax.set_xticks([1975, 1980, 1985, 1990, 1995, 2000, 2005])
ax.set_xticklabels(['1970s','1980','1985','1990','1995','2000','2005'],minor=False)
ax.legend(frameon=True, loc='upper right', fontsize=8, framealpha=1, edgecolor='white')

plt.savefig(figpath + 'model_probs_sens_rho_xi.pdf',bbox_inches='tight')
extract_fig_data(ax,figpath + 'model_probs_sens_rho_xi.csv')
plt.close('all')


# ----------------------------------
# AGS sensitivity

labels=[r'$\omega(2,1)$'] + [r'$\omega_{%d}(1,2)$'%y for y in range(1980,2007)]
J = np.genfromtxt(inpath + 'AGS_jacobian_baseline.txt',delimiter=' ')
J2 = np.genfromtxt(inpath + 'AGS_jacobian_permtpu.txt',delimiter=' ')

#for i in range(9,J.shape[1]):
#    J[:,i] = J[:,i] - J[0,i]
#    J2[:,i] = J2[:,i] - J2[0,i]

years=range(1974,2009)
fig,ax = plt.subplots(5,6,figsize = (9,6),sharex=True,sharey=False)
fig2,ax2 = plt.subplots(5,6,figsize = (9,6),sharex=True,sharey=False)

row=0
col=0
for i in range(28):

    c = [colors[0] for x in range(len(J[:,i]))]
    if i==0:
        for j in range(5):
            c[j] = colors[1]
    else:
        c[i+6] = colors[1]

    ax[row][col].axhline(0,lw=0.2,color='black',alpha=0.5,zorder=0)
    ax2[row][col].axhline(0,lw=0.2,color='black',alpha=0.5,zorder=0)
    
    ax[row][col].vlines(years[:-1],0.0,J[:,i],colors='k',lw=0.2,zorder=1)
    ax2[row][col].vlines(years[:-1],0.0,J2[:,i],colors='k',lw=0.2,zorder=2)
    
    ax[row][col].scatter(years[:-1],J[:,i],color=c,s=1,label='Baseline',zorder=3)
    ax2[row][col].scatter(years[:-1],J2[:,i],color=c,s=1,label='Surprise',zorder=4)

    ax[row][col].set_title(labels[i],size=8,pad=2)
    ax2[row][col].set_title(labels[i],size=8,pad=2)

    col=col+1
    if col>5:
        row = row+1
        col=0

for a in ax.reshape(-1):
    a.set_xticks([])
    a.set_yticks([])    
    a.set_xticklabels([])
    a.set_yticklabels([])

for a in ax[-1,:]:
    a.set_xticks([1975,1980,1985,1990,1995,2000,2005])
    a.set_xticklabels([1975,1980,1985,1990,1995,2000,2005],size=4.5)

for a in ax2.reshape(-1):
    a.set_xticks([])
    a.set_yticks([])    
    a.set_xticklabels([])
    a.set_yticklabels([])

for a in ax2[-1,:]:
    a.set_xticks([1975,1980,1985,1990,1995,2000,2005])
    a.set_xticklabels([1975,1980,1985,1990,1995,2000,2005],size=4.5)
    
fig.tight_layout()
fig2.tight_layout()

plt.figure(fig)
plt.savefig(figpath + 'prob_jacobian_baseline.pdf',bbox_inches='tight')

plt.figure(fig2)
plt.savefig(figpath + 'prob_jacobian_permtpu.pdf',bbox_inches='tight')

plt.close('all')

